import torch
import torch.nn as nn

nc = 3
ndf = 32


# Discriminator
class Discriminator(nn.Module):
    def __init__(self):
        super(Discriminator, self).__init__()
        self.main = nn.Sequential(
            nn.Conv2d(nc, ndf, 4, 2, 1, bias=False),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(ndf, ndf * 2, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ndf * 2), nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(ndf * 2, ndf * 4, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ndf * 4), nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(ndf * 4, ndf * 8, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ndf * 8), nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(ndf * 8, ndf * 16, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ndf * 16), nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(in_channels=ndf * 16,
                      out_channels=1,
                      kernel_size=(12, 4),
                      bias=False), nn.Sigmoid())

    def forward(self, inputs):
        return self.main(inputs)


x = torch.ones((1, 3, 384, 128), dtype=torch.float32)
model = Discriminator()
y = model(x)
print(y.shape)
